//全局限流器
package server

import (
	"net/http"
	"stforum/global"
	"sync"

	"github.com/gin-gonic/gin"
	"golang.org/x/time/rate"
)

type ServerTimeLimiter struct {
	Limiters sync.Map
}

func NewServerTimeLimiter() *ServerTimeLimiter {
	return &ServerTimeLimiter{}
}

func (s *ServerTimeLimiter) AddIP(ip string) {
	s.Limiters.Store(ip, rate.NewLimiter(5, 20))
}

func (s *ServerTimeLimiter) Allow(ip string) bool {
	_, have := s.Limiters.Load(ip)
	if !have {
		s.AddIP(ip)
		return true
	}
	i, _ := s.Limiters.Load(ip)
	res := i.(*rate.Limiter).Allow()
	return res
}

//gin中间件

func (s *ServerTimeLimiter) Limit() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		b := s.Allow(ctx.ClientIP())
		if b {
			global.Logger.Debug("允许ip", ctx.ClientIP(), "通行")
			ctx.Next()
			return
		} else {
			global.Logger.Warn("禁止ip", ctx.ClientIP(), "通行", "访问频率异常！")
			ctx.JSON(http.StatusBadRequest, gin.H{"status": "stop", "msg": "访问频率异常"})
			ctx.Abort()
			return
		}
	}
}
